/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;

import edu.berkeley.nlp.util.Pair;
import edu.umd.clip.lm.util.BinaryPrefix;
import edu.umd.clip.lm.util.SimpleList;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.VariationalTaggerTrellisSlice.Trie;
import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.factors.Dictionary;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public abstract class Tagger {
	final int beamSize;
	final Decoder decoder;
	final BinaryPrefix tags[];
	final byte tagFactorIdx;
	
	protected long taggingTime = 0;
	protected long updatePathsTime = 0;
	protected long makingTaggerSliceTime = 0;

	public Tagger(Decoder decoder, int beamSize, String tagFactor) {
		this.decoder = decoder; 
		this.beamSize = beamSize;
		Experiment exp = Experiment.getInstance();
		this.tagFactorIdx = exp.getTupleDescription().getFactorIndex(tagFactor);
		this.tags = exp.getTags(tagFactorIdx);
		
		TaggerNbestItem.tagFactor = this.tagFactorIdx;
	}

	public abstract Pair<SortedSet<? extends TaggedPath>, Double> evaluateSentence(Collection<FactorTuple> sentence);
	public abstract Pair<SortedSet<? extends TaggedPath>, Double> evaluateSentence(long[] sentence);
	
	/*
	public Pair<SortedSet<? extends TaggedPath>, Double> evaluateSentence(long[] sentence) {
		ArrayList<FactorTuple> tuples = new ArrayList<FactorTuple>(sentence.length);
		for(long tuple : sentence) {
			tuples.add(new FactorTuple(tuple));
		}
		return evaluateSentence(tuples);
	}
	*/
	
	static class TaggerNbestItem implements Comparable<TaggerNbestItem>, TaggedPath {
		final TaggerState state;
		final double prob;
		final SimpleList<Integer> path;
		private static byte tagFactor;

		/* (non-Javadoc)
		 * @see java.lang.Comparable#compareTo(java.lang.Object)
		 */
		@Override
		public int compareTo(TaggerNbestItem o) {
			return (int) Math.signum(prob - o.prob);
		}
		
		public TaggerNbestItem(TaggerState state, double prob, SimpleList<Integer> path) {
			this.state = state;
			this.prob = prob;
			this.path = path;
		}
		
		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.decoding.TaggedPath#getLogProb()
		 */
		public double getLogProb() {
			return prob;
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.decoding.TaggedPath#getOrderedTags()
		 */
		public int[] getOrderedTags() {
			int t[] = new int[path.length()];
			SimpleList<Integer> tail = path;
			for(int i=1; i<=t.length; ++i) {
				// reverse to the proper order
				t[t.length-i] = tail.head().intValue();
				tail = tail.tail();
			}
			return t;
		}
		
		public String toString() {
			StringBuilder sb = new StringBuilder();
			Dictionary dict = Experiment.getInstance().getTupleDescription().getDictionary(tagFactor);
			
			int tags[] = getOrderedTags();
			for(int tag : tags) {
				sb.append(dict.getWord(tag));
				sb.append(' ');
			}
			sb.append('(');
			sb.append(prob);
			sb.append(')');
			return sb.toString();
		}

	}
	
	static class TaggerState {
		int ctx[]; // in order i, i-1, i-2, ...
		double emit_p;
		
		TreeSet<TaggerNbestItem> nbest;

		public TaggerState(int ctx[], double prob) {
			this.ctx = ctx.clone();
			this.emit_p = prob;
			nbest = new TreeSet<TaggerNbestItem>();
		}
		
		public void addProb(double prob) {
			this.emit_p += prob;
		}
		
		public double getProb() {
			return emit_p;
		}
	}

	static class Trie implements Comparable<Trie> {
		Trie subTries[];
		final int tag;
		TaggerState state;
		
		public Trie(int tag, TaggerState state) {
			this.tag = tag;
			this.state = state;
		}

		public Trie(int tag) {
			this(tag, null);
		}
		
		public int compareTo(Trie other) {
			return tag - other.tag;
		}
		
		public Trie getSubtrie(int tag) {
			if (subTries == null) return null;
			for(int i=0; i<subTries.length; ++i) {
				if (subTries[i].tag == tag) {
					return subTries[i];
				}
			}
			return null;
		}
		
		public void addSubtrie(Trie t) {
			if (subTries == null) {
				subTries = new Trie[1];
			} else {
				subTries = Arrays.copyOf(subTries, subTries.length+1);
			}
			subTries[subTries.length-1] = t;
		}
		
		/**
		 * @param ctx the context is organized as follows: t[i], t[i-1], t[i-2], ...
		 * 
		 * @return
		 */
		public Trie findTrie(int ctx[]) {
			return findTrie(ctx, 0);
		}
		
		Trie findTrie(int ctx[], int idx) {
			int t = ctx[idx];
			Trie subTrie = getSubtrie(t);
			if (idx == ctx.length - 1) {
				return subTrie;
			}
			if (subTrie == null) return null;
			return subTrie.findTrie(ctx, idx+1);
		}
		
		/**
		 * @param ctx the context is organized as follows: t[i], t[i-1], t[i-2], ...
		 * @param prob
		 */
		public void addProb(int ctx[], double prob) {
			addProb(ctx, 0, prob);
		}
		
		void addProb(int ctx[], int idx, double prob) {
			//this.prob += prob;
			int tag = ctx[idx];
			
			// search for tag
			Trie subTrie = getSubtrie(tag);
			if (subTrie == null) {
				if (idx == ctx.length - 1) {
					// add item
					Trie t = new Trie(tag, new TaggerState(ctx,prob));
					addSubtrie(t);
					return;
				}
				// add sub trie
				subTrie = new Trie(tag);
				addSubtrie(subTrie);
			} else if (idx == ctx.length - 1) {
				subTrie.state.addProb(prob);
				return;
			}
			subTrie.addProb(ctx, idx+1, prob);
		}

	}
	
	static class TaggerTrellisSlice {
		Trie states;
		
		public TaggerTrellisSlice() {
			this.states = new Trie(-1);
		}
		
		void addProb(int ctx[], double prob) {
			states.addProb(ctx, prob);
		}
		
		TaggerState getState(int ctx[]) {
			Trie trie = states.findTrie(ctx);
			return trie == null ? null : trie.state;
		}
		
		void makeNbest(TaggerTrellisSlice prevSlice, final int N, TreeSet<TaggerNbestItem> overallNbest) {
			// iterate states
			makeNbest(states, prevSlice, N, overallNbest);
		}
		
		void makeNbest(Trie trie, TaggerTrellisSlice prevSlice, final int N, TreeSet<TaggerNbestItem> overallNbest) {
			// iterate states
			if (trie.state != null) {
				makeNbest(trie.state, prevSlice, N);
				
				if (overallNbest != null) {
					double lowProb = Double.NEGATIVE_INFINITY;
					for(Iterator<TaggerNbestItem> it = trie.state.nbest.descendingIterator(); it.hasNext(); ) {
						TaggerNbestItem nbestItem = it.next();
						if (nbestItem.prob < lowProb) break;

						overallNbest.add(new TaggerNbestItem(trie.state, nbestItem.prob, nbestItem.path));
						if (overallNbest.size() > N) {
							TaggerNbestItem lowest = overallNbest.pollFirst();
							lowProb = lowest.prob;
						}
					}					
				}
			} else {
				for(Trie subtrie : trie.subTries) {
					makeNbest(subtrie, prevSlice, N, overallNbest);
				}
			}
		}
		
		void makeNbest(TaggerState state, TaggerTrellisSlice prevSlice, final int N) {
			int prevCtx[] = Arrays.copyOfRange(state.ctx, 1, state.ctx.length);
			
			Trie trie = prevSlice.states.findTrie(prevCtx);
			if (trie == null) {
				System.err.println("oops! can't findTrie in makeNbest");
				return;
			}
			double stateLogProb = Math.log10(state.emit_p);
			
			double lowProb = Double.NEGATIVE_INFINITY;
			TreeSet<TaggerNbestItem> nbest = new TreeSet<TaggerNbestItem>();
			
			for(Trie subtrie : trie.subTries) {
				TaggerState prevState = subtrie.state;
				assert(prevState != null);
				
				for(Iterator<TaggerNbestItem> it = prevState.nbest.descendingIterator(); it.hasNext(); ) {
					TaggerNbestItem nbestItem = it.next();
					
					double prob = nbestItem.prob + stateLogProb;
					if (lowProb > prob) {
						// iterating in descending order -- all others will be lower
						break;
					}
					SimpleList<Integer> path = nbestItem.path.add(Integer.valueOf(state.ctx[0]));
					
					nbest.add(new TaggerNbestItem(prevState, prob, path));
					if (nbest.size() > N) {
						TaggerNbestItem lowest = nbest.pollFirst();
						lowProb = lowest.prob;
					}
				}
			}
			state.nbest = nbest;
		}
	}
	
}
